import json, os
from functools import cached_property
from typing import List, Dict, Union, Optional, Any, Tuple
from config import DATABASE_DIR


class DatabaseSchema():

    def __init__(self, database: str) -> None:
        """ Initialize the database schema object.
        """
        self.database_name = database
        self.database_schema = self._load_database_schema(self.database_name)


    def _load_database_schema(self, database_name: str):
        json_path = os.path.join(DATABASE_DIR, database_name, database_name + '.json')
        if os.path.exists(json_path):
            with open(json_path, 'r') as f:
                schema = json.load(f)
        else:
            raise FileNotFoundError(f"File {json_path} not found.")
        return schema


    @cached_property
    def id2table_mapping(self) -> List[str]:
        """ 输出数据库中所有表名的列表"""
        return [table['table']['table_name'] for table in self.database_schema['database_schema']]


    @cached_property
    def table2id_mapping(self) -> Dict[str, int]:
        """ 获取表名 -> 表 ID 映射字典，ID 根据字段 `database_schema` 中的顺序定义，从 0 开始。"""
        return {table['table']['table_name']: idx for idx, table in enumerate(self.database_schema['database_schema'])}


    @property
    def tables(self) -> List[str]:
        """ 输出数据库中所有表名的列表"""
        return self.id2table_mapping


    def table2id(self, table_name: str) -> int:
        """ 根据表名获取表ID。"""
        return self.table2id_mapping[table_name]


    def id2table(self, table_id: int) -> str:
        """ 根据表ID获取表名。"""
        return self.id2table_mapping[table_id]


    @cached_property
    def table2column_mapping(self) -> Dict[str, List[str]]:
        """ 获取每个表的列名列表映射字典。"""
        return {table['table']['table_name']: [col['column_name'] for col in table['columns']] for table in self.database_schema['database_schema']}


    def table2column(self, table_name: Union[int, str]) -> List[str]:
        """ 根据表名给出表的列名列表。"""
        table_name = self.id2table(table_name) if type(table_name) == int else table_name
        return self.table2column_mapping[table_name]


    def get_primary_keys(self, table_name: Union[int, str]) -> List[str]:
        """ 根据表名获取表的主键列表"""
        table_id = table_name if type(table_name) == int else self.table2id_mapping[table_name]
        return self.database_schema['database_schema'][table_id]['primary_keys']


    def is_encodable(self, table_name: str, column_name: str) -> bool:
        """ 检测指定表的指定列是否可编码（即是否包含可用于向量化的内容）。如果提供了模态参数，则进一步检查该列是否与指定的模态相关联。"""
        for column in self.database_schema['database_schema'][self.table2id_mapping[table_name]]['columns']:
            if column['column_name'] == column_name:
                if column.get('encodable') is not None:
                    return True
        return False


    @cached_property
    def id2column_mapping(self) -> Dict[int, str]:
        """获取列 ID到列名映射字典，ID 根据所有表的 `columns` 字段中的顺序定义，从 0 开始。请注意，列 ID 在数据库中是全局唯一的，而不是在局部表中唯一的。"""
        return [col['column_name'] for table in self.database_schema['database_schema'] for col in table['columns']]

    @cached_property
    def column2id_mapping(self) -> Dict[str, int]:
        """获取列名到列 ID 的映射字典，ID 根据所有表的 `columns` 字段中的顺序定义，从 0 开始。请注意，列 ID 在数据库中是全局唯一的，而不是在局部表中唯一的。"""
        return {col['column_name']: idx for idx, col in enumerate(self.id2column_mapping)}

    def column2id(self, column_name: str) -> int:
        """ 根据列名获取列 ID。"""
        return self.column2id_mapping[column_name]
    
    def id2column(self, column_id: int) -> str:
        """ 根据列 ID 获取列名。"""
        return self.id2column_mapping[column_id]


    # TODO: add more utility methods or properties to get the database schema information, e.g., mapping column_name to its data type, etc.
    pass


if __name__ == '__main__':
    db_schema = DatabaseSchema('DHF_reservoir')
    print("Tables:", db_schema.is_encodable("tables", "table_summary"))